package main

import (
	"crypto/tls"
	"fmt"
	log "github.com/sirupsen/logrus"
	"github.com/urfave/cli"
	"github.com/vbauerster/mpb/v7"
	"github.com/vbauerster/mpb/v7/decor"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"sync"
)

type FileChunk struct {
	start  int64
	end    int64
	status int8
}

// 下载指定分段（更新总进度条）
func processResponse(resp *http.Response, file *os.File, start int64, bar *mpb.Bar) error {
	defer resp.Body.Close()

	// 将分段写入文件
	buffer := make([]byte, 32*1024) // 32KB 缓冲区
	for {
		n, err := resp.Body.Read(buffer)
		if n > 0 {
			_, writeErr := file.WriteAt(buffer[:n], start)
			if writeErr != nil {
				return writeErr
			}
			start += int64(n)
			bar.IncrBy(n) // 更新总进度条
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
	}

	return nil
}

// 全量下载文件
func fullDownload(url string, outFile *os.File) {
	// 发起 GET 请求
	resp, err := http.Get(url)
	if err != nil {
		fmt.Println("Error making GET request:", err)
		return
	}
	defer resp.Body.Close()

	defer outFile.Close()

	// 将响应体写入文件
	_, err = io.Copy(outFile, resp.Body)
	if err != nil {
		fmt.Println("Error writing to file:", err)
		return
	}

	fmt.Println("Full download complete")
}

func run(c *cli.Context) error {
	raw_url := c.String("url")

	parsedUrl, err := url.Parse(raw_url)

	if err != nil {
		log.Fatal("parse url error")
	}

	req, err := http.NewRequest("HEAD", raw_url, nil)
	if err != nil {
		log.Fatal("create request error")
	}
	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, // 跳过证书验证
		},
	}
	resp, err := client.Do(req)

	if err != nil {
		log.Infof("url: %s, err: %s\n", raw_url, err)
		log.Fatal("do request error")
	}
	defer resp.Body.Close()

	file_name := getFileNameFromHeader(resp.Header)

	if file_name == "" {
		file_name = getFileNameFromUrl(parsedUrl)
	}

	file, err := os.Create(file_name)
	defer file.Close()

	content_length := resp.Header.Get("Content-Length")
	acceptRanges := resp.Header.Get("Accept-Ranges")

	// 初始化总进度条
	p := mpb.New(mpb.WithWidth(60))

	if acceptRanges != "bytes" || content_length == "" {
		log.Infof("不支持分段下载，开始全量下载\n")
		fullDownload(raw_url, file)
		return nil
	}
	contentLength, err := strconv.ParseInt(content_length, 10, 64)

	bar := p.AddBar(int64(contentLength),
		mpb.PrependDecorators(
			decor.CountersKibiByte("% .2f / % .2f"),
		),
		mpb.AppendDecorators(
			decor.Percentage(),
		),
	)
	chunkSize := int64(1 * 1024 * 1024)

	// 使用通道保存响应
	responseChan := make(chan *http.Response, 10)
	errorChan := make(chan error, 10)

	// 使用多个协程分段发送 GET 请求
	var requestWg sync.WaitGroup
	for start := int64(0); start < contentLength; start += chunkSize {
		end := start + chunkSize - 1
		if end >= contentLength {
			end = contentLength - 1
		}

		requestWg.Add(1)
		go func(start, end int64) {
			defer requestWg.Done()
			req, err := http.NewRequest("GET", raw_url, nil)
			if err != nil {
				errorChan <- err
				return
			}

			// 设置 Range 请求头
			rangeHeader := fmt.Sprintf("bytes=%d-%d", start, end)
			req.Header.Set("Range", rangeHeader)

			// 发起请求
			client := &http.Client{}
			resp, err := client.Do(req)
			if err != nil {
				errorChan <- err
				return
			}

			// 将响应发送到通道
			responseChan <- resp
		}(start, end)
	}

	// 等待所有请求发送完成
	go func() {
		requestWg.Wait()
		close(responseChan)
		close(errorChan)
	}()

	// 多线程处理响应
	var processWg sync.WaitGroup
	workerCount := 200 // 并发处理响应的 Goroutine 数量
	for i := 0; i < workerCount; i++ {
		processWg.Add(1)
		go func() {
			defer processWg.Done()
			for resp := range responseChan {
				// 获取 Range 起始位置
				contentRange := resp.Header.Get("Content-Range")
				var start int64
				fmt.Sscanf(contentRange, "bytes %d-", &start)

				// 处理响应内容
				err := processResponse(resp, file, start, bar)
				if err != nil {
					log.Errorf("Error processing response: %v", err)
				}
			}
		}()
	}

	// 等待所有响应处理完成
	processWg.Wait()

	// 检查是否有错误
	for err := range errorChan {
		log.Errorf("Error occurred: %v", err)
	}

	p.Wait()

	return nil
}
